#include "kdtree.hpp"
namespace chen {

Kdtree::Kdtree (const std::vector<const Shape *> & v): Structure::Structure(v) {
    int n = v.size();
    LOG(INFO) << "[Kdtree] start building, total: " << n << " objects";
    dtot = 0;
    d = new Node[n+1];
    build(0, n);
}

int cpd;
bool comp (const Shape * a, const Shape * b) {
    return a->wbnd.mx[cpd] < b->wbnd.mx[cpd];
}

int Kdtree::newNode () {
    dtot++;
    d[dtot].l = d[dtot].r = 0;
    return dtot;
}


void Kdtree::query_x (int u, const Ray & r, SurfInteract * si, bool & b) const {
    if(d[u].bnd.interact(r) > r.t) return ;
    if(d[u].s->wbnd.interact(r) <= r.t) b |= d[u].s->interact(r, si);
    Float t1 = Infinity, t2 = Infinity;
    if(d[u].l) t1 = d[d[u].l].bnd.interact(r);
    if(d[u].r) t2 = d[d[u].r].bnd.interact(r);
    if(t1 < t2) {
        if(d[u].l && t1 < r.t) query_x(d[u].l, r, si, b);
        if(d[u].r && t2 < r.t) query_x(d[u].r, r, si, b);
    } else  {
        if(d[u].r && t2 < r.t) query_x(d[u].r, r, si, b);
        if(d[u].l && t1 < r.t) query_x(d[u].l, r, si, b);
    }
}
const Shape * Kdtree::query_y (int u, const Ray & r) const {
    if(d[u].bnd.interact(r) > r.t) return nullptr;
    const Shape * ret = nullptr, * tmp;
    if(d[u].s->wbnd.interact(r) <= r.t && d[u].s->mask(r)) ret = d[u].s;
    Float t1 = Infinity, t2 = Infinity;
    if(d[u].l) t1 = d[d[u].l].bnd.interact(r);
    if(d[u].r) t2 = d[d[u].r].bnd.interact(r);
    if(t1 < t2) {
        if(d[u].l && t1 < r.t) if((tmp = query_y(d[u].l, r))) ret = tmp;
        if(d[u].r && t2 < r.t) if((tmp = query_y(d[u].r, r))) ret = tmp;
    } else  {
        if(d[u].r && t2 < r.t) if((tmp = query_y(d[u].r, r))) ret = tmp;
        if(d[u].l && t1 < r.t) if((tmp = query_y(d[u].l, r))) ret = tmp;
    }
    return ret;
}

int Kdtree::build (int l, int r) {
    if(l >= r) return 0;
    int u = newNode();
    if(l+1 == r) {
        d[u].s = shapes[l];
        d[u].bnd = shapes[l]->wbnd;
        return u;
    }
    Bound3 bnd = shapes[l]->wbnd;
    for(int i = l+1; i<r; i++) bnd = combine(bnd, shapes[i]->wbnd);
    cpd = 0;
    for(int i = 0; i<3; i++) 
        if(bnd.mx[i]-bnd.mn[i] > bnd.mx[cpd]-bnd.mn[cpd]) cpd = i;
    int mid = (l+r)>>1;
    std::nth_element(shapes.begin()+l, shapes.begin()+mid, shapes.begin()+r, comp);
    d[u].s = shapes[mid];
    d[u].bnd = bnd;
    d[u].l = build(l, mid);
    d[u].r = build(mid+1, r);
    return u;
}


bool Kdtree::castRay (const Ray & r, SurfInteract * si) const {
    bool b = false;
    query_x(1, r, si, b);
    return b;
}

const Shape * Kdtree::testRay (const Ray & r) const {
    return query_y(1, r);
}

Kdtree::~Kdtree () {
    delete [] d;
}

}